
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@NAME          : user_dto.py
@TIME          : 2024/11/16 21:24:58
@AUTHOR        : chenlip
@VERSION       : 0.0.1
@DESCRIPTION   : 用户模型的CRUD数据传输对象的基础类
'''

import base64
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from flask import current_app

from sqlalchemy.exc import SQLAlchemyError
from init_app import sqldb as db

class BaseDTO:
    """基础DTO类，提供通用方法"""
    
    @staticmethod
    def handle_db_error(e: SQLAlchemyError, rollback: bool = True) -> dict:
        """处理数据库异常"""
        if rollback:
            db.session.rollback()
        error_msg = str(e)

        return {"success": False, "message": f"数据库操作失败: {error_msg}"}
    
    @staticmethod
    def aes_encrypt(data: str, key: str=None) -> str:
        """
        AES256加密
        
        Args:
            data: 要加密的字符串
            key: 加密密钥
            
        Returns:
            加密后的Base64编码字符串
        """
        if not data:
            return ""
            
        try:
            # 确保数据是字节类型
            if isinstance(data, str):
                data = data.encode('utf-8')
                
            # 确保密钥是32字节(256位)
            if not key:
                key = current_app.config.get('SECRET_KEY', 'hoeking#1975@1234567890')

            key_bytes = key.encode('utf-8')
            key_bytes = key_bytes.ljust(32, b'\0')[:32]  # 填充或截断至32字节
            
            # 生成随机初始向量
            iv = os.urandom(16)
            
            # 创建AES-CBC模式的加密器
            cipher = AES.new(key_bytes, AES.MODE_CBC, iv)
            
            # 对数据进行填充并加密
            encrypted_data = cipher.encrypt(pad(data, AES.block_size))
            
            # 将IV和加密数据组合并Base64编码
            combined = iv + encrypted_data
            result = base64.b64encode(combined).decode('utf-8')
            
            return result
        except Exception as e:
            # 记录错误
            if hasattr(current_app, 'logger') and current_app.logger:
                current_app.logger.error(f"AES加密错误: {str(e)}")
            else:
                print(f"AES加密错误: {str(e)}")
            return ""

    @staticmethod
    def aes_decrypt(encrypted_data: str, key: str=None) -> str:
        """
        AES256解密
        
        Args:
            encrypted_data: 加密的Base64编码字符串
            key: 解密密钥
            
        Returns:
            解密后的原始字符串
        """
        if not encrypted_data:
            return ""
            
        try:
            # Base64解码
            encrypted_bytes = base64.b64decode(encrypted_data)
            
            # 确保密钥是32字节(256位)
            if not key:
                key = current_app.config.get('SECRET_KEY', 'hoeking#1975@1234567890')

            key_bytes = key.encode('utf-8')
            key_bytes = key_bytes.ljust(32, b'\0')[:32]  # 填充或截断至32字节
            
            # 提取IV（前16字节）
            iv = encrypted_bytes[:16]
            actual_encrypted_data = encrypted_bytes[16:]
            
            # 创建AES-CBC模式的解密器
            cipher = AES.new(key_bytes, AES.MODE_CBC, iv)
            
            # 解密并去除填充
            decrypted_data = unpad(cipher.decrypt(actual_encrypted_data), AES.block_size)
            
            # 转换为字符串
            return decrypted_data.decode('utf-8')
        except Exception as e:
            # 记录错误
            if hasattr(current_app, 'logger') and current_app.logger:
                current_app.logger.error(f"AES解密错误: {str(e)}")
            else:
                print(f"AES解密错误: {str(e)}")
            return ""